import random
# from gmssl import func
import string
from gmssl.sm4 import CryptSM4, SM4_ENCRYPT, SM4_DECRYPT

crypt_sm4 = CryptSM4()


def gen_key(bits):
    # key = random.randbytes(15)
    # key = random.sample('zyxwvutsrqponmlkjihgfedcba1234567890', 5)
    num_set = [chr(i) for i in range(48, 58)]
    char_set = [chr(i) for i in range(97, 123)]
    total_set = num_set + char_set
    key = "".join(random.sample(total_set, bits))
    # key = ''.join(random.sample(string.ascii_letters + string.digits, 15))
    return key


def iv_gen(bits):
    num_set = [chr(i) for i in range(48, 58)]
    char_set = [chr(i) for i in range(97, 123)]
    total_set = num_set + char_set
    iv = "".join(random.sample(total_set, bits))
    # iv = ''.join(random.sample(string.ascii_letters + string.digits, 15))
    return iv


def SM4_ecb_enc(plaintext, key):
    crypt_sm4.set_key(key, SM4_ENCRYPT)
    ciphertext_bytes = crypt_sm4.crypt_ecb(plaintext)

    return ciphertext_bytes


def SM4_ecb_dec(ciphertext, key):
    crypt_sm4 = CryptSM4()
    crypt_sm4.set_key(key, SM4_DECRYPT)
    plaintext_bytes = crypt_sm4.crypt_ecb(ciphertext)

    return plaintext_bytes


def SM4_cbc_enc(plaintext, iv, key):
    crypt_sm4.set_key(key, SM4_ENCRYPT)
    ciphertext_bytes = crypt_sm4.crypt_cbc(iv, plaintext)  # bytes类型

    return ciphertext_bytes


def SM4_cbc_dec(ciphertext, iv, key):
    crypt_sm4.set_key(key, SM4_DECRYPT)
    plaintext_bytes = crypt_sm4.crypt_cbc(iv, ciphertext)

    return plaintext_bytes


def exp_SM4_str(plaintext):
    plaintext_bytes = bytes(plaintext, encoding="utf8")
    return plaintext_bytes


def exp_SM4_file(filename):
    with open(filename, "rb") as f:
        plaintext_bytes = f.read()
    f.close()
    return plaintext_bytes


if __name__ == '__main__':

    key = gen_key(16)
    print(type(key))
    print(key)
    key_bytes = bytes(key, encoding="utf8")
    print(len(key_bytes.hex()))
    with open('D:/key_bytes.txt', 'wb') as f:
        f.truncate()
        f.write(key_bytes)
    f.close
    print("key为：", key_bytes.hex())

    iv = iv_gen(16)
    iv_bytes = bytes(iv, encoding="utf8")
    with open('D:/iv_bytes.txt', 'wb') as f:
        f.truncate()
        f.write(iv_bytes)
    f.close
    print("iv为：", iv_bytes.hex())

    choice = input("选择字符串加密还是文件加密")
    if (choice == 'file'):
        print("请输入需要加密的文件名：")
        plaintext_bytes = exp_SM4_file(input())

    else:
        print("请输入需要加密的内容：")
        plaintext_bytes = exp_SM4_str(input())

    ciphertext = SM4_ecb_enc(plaintext_bytes, key_bytes)
    print(type(plaintext_bytes))
    with open('D:/ciphertext_ecb.txt', 'wb') as f:
        f.truncate()
        f.write(ciphertext)
    f.close
    print("ecb加密为：", ciphertext)

    plaintext = SM4_ecb_dec(ciphertext, key_bytes)
    with open('D:/plaintext_ecb.txt', 'wb') as f:
        f.truncate()
        f.write(plaintext)
    f.close
    print("ecb解密为：", plaintext)

    ciphertext = SM4_cbc_enc(plaintext, iv_bytes, key_bytes)
    with open('D:/ciphertext_cbc.txt', 'wb') as f:
        f.truncate()
        f.write(ciphertext)
    f.close
    print("cbc加密为：", ciphertext)

    plaintext = SM4_cbc_dec(ciphertext, iv_bytes, key_bytes)
    with open('D:/plaintext_cbc.txt', 'wb') as f:
        f.truncate()
        f.write(plaintext)
    f.close
    print("cbc解密为：", plaintext)
